import pandas as pd


# 定义含沙量填补函数
def fill_sediment(row, df, sediment_column):
    if pd.isna(row[sediment_column]):
        index = df.index.tolist().index(row.name)
        d_forward = 1
        d_backward = 1
        max_index = len(df) - 1
        min_index = 0

        # 查找下一个非空值
        next_non_na_index = None
        while index + d_forward <= max_index:
            if not pd.isna(df.iloc[index + d_forward][sediment_column]):
                next_non_na_index = index + d_forward
                break
            d_forward += 1

        # 查找前一个非空值
        prev_non_na_index = None
        while index - d_backward >= min_index:
            if not pd.isna(df.iloc[index - d_backward][sediment_column]):
                prev_non_na_index = index - d_backward
                break
            d_backward += 1

        if next_non_na_index is not None and prev_non_na_index is not None:
            q_st_prev = df.iloc[prev_non_na_index][sediment_column]
            q_st_next = df.iloc[next_non_na_index][sediment_column]
            return (d_forward / (d_forward + d_backward)) * q_st_prev + (
                        d_backward / (d_forward + d_backward)) * q_st_next
        elif prev_non_na_index is not None:
            return df.iloc[prev_non_na_index][sediment_column]
        elif next_non_na_index is not None:
            return df.iloc[next_non_na_index][sediment_column]
    return row[sediment_column]


def calculate_sediment_table():
    target_years = [2016, 2017, 2018, 2019, 2020, 2021]
    file_path = 'ori_message.xlsx'
    try:
        df_all = pd.read_excel(file_path, sheet_name=[str(year) for year in target_years])
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
        return
    except Exception as e:
        print(f"错误：读取文件时发生未知错误：{e}")
        return

    year_sediment = {}

    for year in target_years:
        df = df_all[str(year)]
        # 去除列名中的不可见字符
        df.columns = [col.strip() for col in df.columns]

        # 填充年、月、日列
        df[['年', '月', '日']] = df[['年', '月', '日']].ffill()

        # 筛选8:00数据
        df_8am = df[df['时间'] == '8:00'].copy()

        # 转换日期类型
        df_8am[['年', '月', '日']] = df_8am[['年', '月', '日']].astype(int)
        df_8am.drop(['时间'], axis=1, inplace=True)

        # 生成完整日期列表（处理闰年）
        all_dates = []
        for month in range(1, 13):
            start_date = pd.Timestamp(year=year, month=month, day=1)
            last_day = start_date.days_in_month
            for day in range(1, last_day + 1):
                all_dates.append((year, month, day))
        all_dates_df = pd.DataFrame(all_dates, columns=['年', '月', '日'])
        all_dates_df[['年', '月', '日']] = all_dates_df[['年', '月', '日']].astype(int)

        # 左合并数据
        df_merged = pd.merge(all_dates_df, df_8am, on=['年', '月', '日'], how='left')
        df_merged = df_merged[~((df_merged['月'] == 2) & (df_merged['日'] == 29))]  # 剔除闰年2月29日

        # 检查含沙量列是否存在
        sediment_column = '含沙量(kg/m3)'
        if sediment_column not in df_merged.columns:
            print(f"警告：{year}年数据中不存在列名 '{sediment_column}'，请检查列名。")
            print(f"{year}年数据的列名有：{', '.join(df_merged.columns)}")
            continue

        # 检查流量列是否存在
        flow_column = '流量(m3/s)'
        if flow_column not in df_merged.columns:
            print(f"警告：{year}年数据中缺少流量列，无法计算！")
            continue

        # 填补含沙量
        df_merged[sediment_column] = df_merged.apply(lambda row: fill_sediment(row, df_merged, sediment_column), axis=1)

        # 计算年总排沙量
        df_merged['每日排沙量(kg)'] = (
                df_merged[sediment_column] *  # 含沙量
                df_merged[flow_column] *  # 流量
                86400  # 每日秒数
        )
        total_kg = df_merged['每日排沙量(kg)'].sum()
        total_billion_ton = total_kg * 1e-11  # 转换为亿吨
        year_sediment[year] = round(total_billion_ton, 3)

    # 输出表4格式结果
    print("\n2016 - 2021年的年总排沙量表")
    print("| 日期   | 2016  | 2017  | 2018  | 2019  | 2020  | 2021  |")
    print("|--------|" + "|".join([" {:.3f} ".format(year_sediment[y]) for y in target_years]) + "|")
    print("单位：亿吨")


# 执行函数
calculate_sediment_table()
